from typing import List
from pathlib import Path

from transformers import BertConfig
from docx import Document

from pluginUtils import BasicReaderPlugin
from pymodels.models.dpr import tinybertVectorizer
from pymodels.utils.BookReader import BasicBookReader
from pymodels.utils.Indexer import FaissIndexerHelper


class ReaderPlugin(BasicReaderPlugin, BasicBookReader):
    def __init__(self, size: int = 512, device: str = "cpu"):
        """
        设置PDF阅读器。
        :param size: 读取文本时的每页的大小。越小，推理器的结果可能越精确。推荐值：512.
        :param device: 推理时使用的处理器。
        :return:
        """
        self.size = size
        self.model_name = "zerohell/tinydpr-acc_0.315-bs_307"
        self.device = device
        dim = BertConfig.from_pretrained(self.model_name).hidden_size
        indexer = FaissIndexerHelper(dim=dim)
        BasicBookReader.__init__(self, indexer=indexer, booktype='.docx', cache=True, arduino=True)
        BasicReaderPlugin.__init__(self)

    def hibernate(self):
        BasicBookReader.del_vectorizer(self)

    def wake(self):
        dpr = tinybertVectorizer(self.model_name)
        dpr.model.to(self.device)
        BasicBookReader.set_vectorizer(self, dpr)

    def file2pages(self, filename: str) -> List[str]:  # noqa
        """
        :param filename: 要处理的文件名称。
        :return: 页形成的数组。会包括一个空页，用于判断是否为不相关的问题。
        """
        size = self.size
        overlap = size // 2
        filename = Path(filename)
        self.readbook_name = filename.name
        doc = Document(filename)
        content = ""
        for para in doc.paragraphs:
            content += para.text
        blocks = [content[i: i + size] for i in range(0, len(content), overlap)]
        return blocks

    def get_page_size(self) -> int:
        return self.size

    def search(self, query: str, k: int):
        return BasicBookReader.search(self, query, k=k, return_page_content=True, return_cos_similarity=True)

    def readbook(self, bookpath: str):
        return BasicBookReader.readbook(self, filename=bookpath)

    def get_readfile_name(self):
        return BasicBookReader.get_readfile_name(self)

    def support_file_type(self):
        return {".docx"}

    def get_introduction(self):
        return """
        作者：xzh
        日期：2023年5月27日
        邮箱：kodderopert@163.com
        """


if __name__ == "__main__":
    from utils.idle import idle
    plugin = ReaderPlugin()
    plugin.readbook(r"D:\dell\desktop\毕业设计\终期答辩\1.docx")
    idle(globals())
